from typing import Any, Dict
from database import DatabaseManager
from config import DEFAULT_DB_CONFIG, logger
from sql_builder import SQLBuilder
from query_parser import QueryParser
import json
from llm_client import LLMClient


class MYSQLMCPServer:
    def __init__(self):
        self.db_manager = DatabaseManager()

    async def connect_database(self, config: Dict[str, Any]):
        return await self.db_manager.connect_database(config)

    async def execute_sql(self, sql: str) -> Dict[str, Any]:
        if not self.db_manager.db_connection:
            await self.connect_database(DEFAULT_DB_CONFIG)
        result = await self.db_manager.execute_sql(sql)
        # 如果执行SQL语句发生了错误，则调用大模型帮助分析错误原因以及如何修改
        if not result.get("success") and "error" in result:
            error_analysis = SQLBuilder.analyze_sql_error(result["error"], sql)
            result["error_analysis"] = error_analysis
        return result

    async def natural_language_query(self, query: str) -> Dict[str, Any]:
        try:
            # 使用QueryParser解析自然语言查询
            parsed = QueryParser.parse_natural_language(query)
            if not parsed:
                return {
                    "success": False,
                    "message": f"无法解析查询意图",
                    "original_query": query,
                }
            logger.error(f"解析结果:{parsed}")
            sql = SQLBuilder.build_sql_query(parsed)
            if not sql:
                return {
                    "success": False,
                    "message": f"无法生成SQL语句",
                    "original_query": query,
                    "parsed": parsed,
                }
            logger.error(f"生成SQL:{sql}")
            result = await self.execute_sql(sql)
            logger.error(f"执行SQL成功:{result}")
            if result.get("success") and "data" in result:
                result["explanation"] = self._explain_query_result(
                    query, result, parsed
                )

            return {
                "success": True,
                "original_query": query,
                "parsed": parsed,
                "generated_sql": sql,
                "execution_result": result,
            }

        except Exception as e:
            logger.error(f"自然语言查询失败:{str(e)}")
            return {
                "success": False,
                "message": f"自然语言查询处理失败:{str(e)}",
                "original_query": query,
            }

    def _explain_query_result(
        self, original_query: str, result: Dict[str, Any], parsed: Dict[str, Any]
    ) -> str:
        try:
            data_summary = ""
            if result.get("data"):
                data_count = len(result["data"])
                if data_count > 0:
                    sample_data = result["data"][:3]
                    data_summary = f"查询到了{data_count}条记录，示例数据:{json.dumps(sample_data,ensure_ascii=False)}"
                else:
                    data_summary = "未查询到任何记录"
            prompt = f"""
                请根据以下信息，用简明中文解释查询结果：
                原始查询: {original_query}
                查询意图: {parsed.get('operation', 'SELECT')} 操作，涉及表: {', '.join(parsed.get('tables', []))}
                查询结果: {data_summary}
                请给出：
                1. 查询结果的简要总结
                2. 若结果为空，可能的原因
                3. 如有需要，给出进一步查询建议
                答案不超过100字。
            """
            explanation = LLMClient.call_llm(prompt)
            if explanation:
                return explanation.strip()
            else:
                return "查询完成"
        except Exception as e:
            logger.warning(f"结果解析失败:{str(e)}")
            return "查询完成"

    def close(self) -> None:
        self.db_manager.close()
